# -*- coding: utf-8 -*-
from collections import Counter
from itertools import product

import torch 
import torch.nn.functional as F

class PathTree(object):
    def __init__(self,root,matrix,device) -> None:
        self.root = root
        self.__build(matrix,device)
        self.path = []

    def __build(self,matrix,device):
        index = torch.arange(len(matrix),device=device)
        self.tree = {
            node:index[matrix[node]].tolist()
            for node in range(len(matrix))
        }
    
    def __get_path(self,pred_path):
        tail = pred_path[-1]
        for next in self.tree[tail]:
            self.__get_path(pred_path + [next])
        self.path.append(pred_path)

    def get_all_path(self):
        self.path = []
        self.__get_path([self.root])
        return self.path[:-1]

class MyGraph:
    
    def __init__(self,keys:list=None,matrix:torch.Tensor=None,device='cuda') -> None:
        if isinstance(keys,list):
            self.keys = keys
        elif keys is None:
            self.keys = []
        else:
            raise ValueError("Error keys type")
        
        self.device = device
        self.key2index = {k:v for v,k in enumerate(self.keys)}

        if matrix is None:
            self.matrix = torch.zeros((len(self.keys), len(self.keys)),device=device)
        elif self.__check_format(matrix):
            self.matrix = torch.tensor(matrix,device=device)
        else:
            raise ValueError("Error table format")
  
    def __check_format(self,matrix:torch.Tensor) -> bool:
        if matrix is None:
            return True
        
        checker = (
            isinstance(matrix,torch.Tensor) and
            matrix.ndim == 2 and
            matrix.shape[0] == matrix.shape[1] == len(self.keys)
        )
        return checker
    
    @property
    def indexs(self):
        return torch.arange(0,len(self.keys),device=self.device)

    def print_graph(self) -> None:
        for k in self.keys:
            index = self.key2index[k]
            message = f"{str(k)} -> " + ", ".join(
                map(lambda i: str(self.keys[i]),
                    (i for i,v in enumerate(self.matrix[index,:]) if v>0))
            )
            print(message)

    #region get basic info
    
    def get_nodes(self):
        return self.keys

    @property
    def nodes(self):
        return self.keys

    def get_edges(self): 
        '''
        Returns a list of tuples (origin, destination, weight)
        '''
        edges = []
        for k in self.keys:
            index = self.key2index[k]
            edges.extend([
                (k,self.keys[i],v) for i,v in enumerate(self.matrix[index,:]) if v>0
            ])
        return edges
    
    def edge_weight(self,o,d):
        return self.matrix[self.key2index[o],self.key2index[d]]

    @property
    def edges(self):
        return self.get_edges()
    
    @property
    def size(self):
        '''
        Returns (number of nodes, number of edges)
        '''
        return len(self.keys), len(self.edges)
    
    #endregion

    #region add nodes and edges    
    
    def add_vertex(self, *vertex):
        for v in vertex:
            if v not in self.key2index:
                self.keys.append(v)
                self.key2index[v] = len(self.keys)-1
                self.matrix = F.pad(self.matrix, (0, 1, 0, 1), mode='constant', value=0)
        
    def add_edge(self, o, d, w:float=1.0, undirect=False):
        self.add_vertex(o,d)
        self.matrix[self.key2index[o],self.key2index[d]] = w
        if undirect:
            self.matrix[self.key2index[d],self.key2index[o]] = w
    
    #endregion

    #region successors, predecessors, adjacent nodes
        
    def get_successors(self, v):
        indexs = self.indexs[self.matrix[self.key2index[v],:] > 0]
        return [self.keys[i] for i in indexs]
    
    def get_predecessors(self, v):
        indexs = self.indexs[self.matrix[:,self.key2index[v]] > 0]
        return [self.keys[i] for i in indexs]
    
    def get_adjacents(self, v):
        is_suc = self.matrix[self.key2index[v],:] > 0
        is_pre = self.matrix[:,self.key2index[v]] > 0
        indexs = self.indexs[torch.logical_or(is_suc,is_pre)]
        return [self.keys[i] for i in indexs]
    
    #endregion
        
    #region degrees

    def out_degree(self, v):
        return torch.count_nonzero(self.matrix[self.key2index[v],:])
    
    def in_degree(self, v):
        return torch.count_nonzero(self.matrix[:,self.key2index[v]])
        
    def degree(self, v):
        return len(self.get_adjacents(v))
        
    def all_degrees(self, deg_type = "inout"):
        '''
        Computes the degree (of a given type) for all nodes.
        deg_type can be "in", "out", or "inout"
        '''
        if deg_type not in ("in","out","inout"):
            raise ValueError("deg_type should be in, out, or inout")
        
        checker = torch.zeros_like(self.matrix)
        if "in" in deg_type:
            checker += self.matrix.T
        if "out" in deg_type:
            checker += self.matrix
        result = {
            key : torch.count_nonzero(checker[i,:])
            for i,key in enumerate(self.keys)
            }
        return result
    
    def highest_degrees(self, all_deg=None, deg_type = "inout", top=10):
        if all_deg is None: 
            all_deg = self.all_degrees(deg_type)
        nodes = Counter(all_deg).most_common(top)
        return [n for n,_ in nodes]
    
    #endregion

    #region topological metrics over degrees

    def mean_degree(self, deg_type="inout"):
        degs = self.all_degrees(deg_type)
        return sum(degs.values()) / float(len(self.keys))
        
    def prob_degree(self, deg_type = "inout"):
        degs = self.all_degrees(deg_type)
        result = {
            k : v/len(self.keys)
            for k,v in Counter(degs.values()).items()
            }
        return result
    
    #endregion
    
    #region BFS and DFS searches
    def __dist_matrix(self,weight=False,diag=0) -> torch.Tensor:
        dist_mat = torch.full_like(self.matrix,torch.inf,device=self.device)
        dist_mat[self.matrix>0] = self.matrix[self.matrix>0]
        if not weight:
            dist_mat[dist_mat != torch.inf] = 1
        dist_mat.fill_diagonal_(diag)
        return dist_mat


    def __bfs(self, s, weight=False) -> torch.Tensor:
        dist_mat = self.__dist_matrix(weight)
        dist = torch.full((len(self.keys),),
                          torch.inf,device=self.device)
        dist[self.key2index[s]] = 0

        while True:
            step = dist.reshape(-1,1) + dist_mat
            dist_step,_ = torch.min(step,axis=0)
            l = self.indexs[dist != dist_step]
            dist = dist_step
            if len(l) == 0:
                break

        return dist
    
    def reachable_bfs(self, v):
        dist = self.__bfs(v)
        result = [
            self.keys[i]
            for i,d in enumerate(dist)
            if d != torch.inf and i != self.key2index[v]
        ]
        return result
        
    def reachable_dfs(self, v):
        l = [v]
        res = []
        while len(l) > 0:
            node = l.pop(0)
            if node != v:
                res.append(node)
            s = 0
            for elem in self.get_successors(node):
                if elem not in res and elem not in l:
                    l.insert(s, elem)
                    s += 1
        return res

    def distance(self, s, d, weight=False):
        dist = self.__bfs(s, weight)
        if dist[self.key2index[d]] == torch.inf:
            return None
        return dist[self.key2index[d]]
    
    def short_path_tree(self,start_node,weight=False) -> torch.Tensor:
        """
        Short Path Tree that build from start node
        Argument:
            weight: (bool) whether to use edge weights for distance calculation.
        Return:
            An numpy ndarray where True means the edge is in Short Path Tree.
        wiki url: https://en.wikipedia.org/wiki/Shortest-path_tree
        """
        if start_node not in self.nodes:
            raise ValueError("Error node")

        dist = self.__bfs(start_node, weight)
        dist[dist == torch.inf] = torch.nan
        in_path = dist.reshape(1,-1) - dist.reshape(-1,1)
        in_path = self.__dist_matrix(weight,diag=torch.inf) == in_path
        return in_path
    
    def shortest_path(self, s, d, weight=False):
        if s == d:
            return 0
        dist = self.__bfs(s, weight)
        if dist[self.key2index[d]] == torch.inf:
            return None
        dist[dist == torch.inf] = torch.nan
        in_path = dist.reshape(1,-1) - dist.reshape(-1,1)
        in_path = self.__dist_matrix(weight,diag=torch.inf) == in_path
        current_node = self.key2index[d]
        path = []
        while True:
            path.append(self.keys[current_node])
            pred = self.indexs[in_path[:,current_node]][0]
            if pred == self.key2index[s]:
                path.append(s)
                break
            current_node = pred
        return path[::-1]

        
    def reachable_with_dist(self, s, weight=False):
        dist = self.__bfs(s,weight)
        result = [
            (self.keys[i],d) 
            for i,d in enumerate(dist)
            if d != torch.inf and i != self.key2index[s]
            ]
        return result
    
    #endregion

    #region mean distances ignoring unreachable nodes

    def mean_distances(self, weight=False):
        tot = 0
        num_reachable = 0
        for k in self.keys: 
            distsk = self.reachable_with_dist(k, weight)
            tot += sum((d for _,d in distsk))
            num_reachable += len(distsk)
        mean_dist = tot / num_reachable
        n = len(self.get_nodes())
        return mean_dist, float(num_reachable)/((n-1)*n)
    
    def closeness_centrality(self, node, weight=False):
        dist = self.reachable_with_dist(node, weight)
        if len(dist) == 0:
            return 0.0
        else:
            return len(dist) / sum((d for _,d in dist))
    
    def all_closeness_centrality(self,weight=False):
        cs = {
            k : self.closeness_centrality(k,weight)
            for k in self.keys
            }
        return cs
    
    def highest_closeness(self, top = 10, weight=False):
        center = Counter({
            node:self.closeness_centrality(node, weight)
            for node in self.keys}
            ).most_common(top)
        return [k for k,_ in center]
    
    def betweenness_centrality(self, node, weight=False):
        total_sp = 0
        sps_with_node = 0
        for s,t in product(self.keys,self.keys):
            if s != t and s != node and t != node:
                sp = self.shortest_path(s, t, weight)
                if sp is not None:
                    total_sp += 1
                    if node in sp:
                        sps_with_node += 1 
        return sps_with_node / total_sp

    def all_betweenness_centrality(self, weight=False, normal=True):
        """
        Use short path tree to reduce algorithmic complexity
        """
        betweenness = torch.zeros(len(self.nodes),device=self.device)

        for start_node in self.nodes:
            start = self.key2index[start_node]
            edge_mat = self.short_path_tree(start_node,weight)
            SPT = PathTree(start,edge_mat,self.device)

            temp = torch.zeros_like(self.matrix,device=self.device)
            num_short_to_leaf = torch.zeros(len(self.nodes),device=self.device)
            for path in SPT.get_all_path():
                leaf = path[-1]
                num_short_to_leaf[leaf] += 1
                temp[leaf,path[1:-1]] += 1
            num_short_to_leaf[num_short_to_leaf==0] = 1
            betweenness += torch.sum(temp/num_short_to_leaf.reshape(1,-1),axis=0)

        C = (len(self.keys)-1)*(len(self.keys)-2) if normal else 1
        result = {
            self.keys[i]:
            float(betweenness[i])/C
            for i in range(len(self.keys))
            }
            
        return result                    
    #endregion
    
    #region cycles

    def node_has_cycle(self, v) -> bool:
        l = [v]
        res = False
        visited = [v]
        while len(l) > 0:
            node = l.pop(0)
            for elem in self.get_successors(node):
                if elem == v:
                    return True
                if elem not in visited: 
                    l.append(elem)
                    visited.append(elem)
        return res
    
    def has_cycle(self) -> bool:
        res = False
        for v in self.keys:
            if self.node_has_cycle(v):
                res = True
                break
        return res
    
    #endregion

    #region clustering
        
    def clustering_coef(self, v):
        is_suc = self.matrix[self.key2index[v],:] > 0
        is_pre = self.matrix[:,self.key2index[v]] > 0
        is_link = torch.logical_or(is_suc,is_pre)
        num_link = torch.sum(is_link)
        if num_link <= 1 :
            return 0
        else:
            num_edge = torch.count_nonzero(self.matrix[is_link][:,is_link])
            return num_edge/(num_link*(num_link-1))
        
    def all_clustering_coefs(self):
        ccs = {
            k : self.clustering_coef(k)
            for k in self.keys
            }
        return ccs
        
    def mean_clustering_coef(self):
        ccs = self.all_clustering_coefs()
        return sum(ccs.values()) / float(len(ccs))
            
    def mean_clustering_perdegree(self, deg_type = "inout"):
        degs = self.all_degrees(deg_type)
        ccs = self.all_clustering_coefs()
        degs_k = Counter(degs.values())
        result = dict()
        for node,cc in ccs.items():
            deg = degs[node]
            result[deg] = result.get(deg,0) + cc/degs_k[deg]
        return result
    
    #endregion


def test_AllB():
    N = 1000
    mat = torch.bernoulli(torch.full((N, N), 0.1))
    mat.fill_diagonal_(0)
    G = MyGraph([str(i) for i in range(N)],mat,device="cpu")
    G.all_betweenness_centrality()


if __name__ == "__main__":
    test_AllB()